#include <linux/init.h>
#include <linux/module.h>
#include <linux/fs.h>
#include <linux/mm.h>
#include <linux/uaccess.h>
#include <linux/slab.h>
#include <linux/vmalloc.h>
#include <linux/device.h>

#define DEVICE_NAME "contig_mmap"
#define PAGE_COUNT 4
#define BUFFER_SIZE (PAGE_COUNT * PAGE_SIZE)

static int major;
static void *buffer;
static unsigned long buffer_phys;
static struct page **pages;
static struct class *cls;

static int contig_mmap_open(struct inode *inode, struct file *filp) {
    printk(KERN_INFO "Contig mmap device opened\n");
    return 0;
}

static int contig_mmap_release(struct inode *inode, struct file *filp) {
    printk(KERN_INFO "Contig mmap device closed\n");
    return 0;
}

static int contig_mmap(struct file *filp, struct vm_area_struct *vma) {
    unsigned long size = vma->vm_end - vma->vm_start;
    unsigned long offset = vma->vm_pgoff << PAGE_SHIFT;
    int i = 0;

    if (size > BUFFER_SIZE)
        return -EINVAL;

    vm_flags_set(vma, VM_DONTEXPAND | VM_DONTDUMP);

    // for (i = 0; i < PAGE_COUNT; i++) {
        if (remap_pfn_range(vma,
                           vma->vm_start + i * PAGE_SIZE,
                           page_to_pfn(pages[i]),
                           PAGE_SIZE*PAGE_COUNT,
                           vma->vm_page_prot))
            return -EAGAIN;
    // }

    printk(KERN_INFO "Memory mapped successfully\n");
    return 0;
}

static const struct file_operations fops = {
    .owner = THIS_MODULE,
    .open = contig_mmap_open,
    .release = contig_mmap_release,
    .mmap = contig_mmap,
};

static int __init contig_mmap_init(void) {
    int i;

    major = register_chrdev(0, DEVICE_NAME, &fops);
    if (major < 0) {
        printk(KERN_ALERT "Failed to register device\n");
        return major;
    }

    cls = class_create(DEVICE_NAME);
    if (IS_ERR(cls)) {
        unregister_chrdev(major, DEVICE_NAME);
        printk(KERN_ALERT "Failed to create class\n");
        return PTR_ERR(cls);
    }

    if (device_create(cls, NULL, MKDEV(major, 0), NULL, DEVICE_NAME) == NULL) {
        class_destroy(cls);
        unregister_chrdev(major, DEVICE_NAME);
        printk(KERN_ALERT "Failed to create device\n");
        return -1;
    }

    pages = kmalloc_array(PAGE_COUNT, sizeof(struct page *), GFP_KERNEL);
    if (!pages) {
        device_destroy(cls, MKDEV(major, 0));
        class_destroy(cls);
        unregister_chrdev(major, DEVICE_NAME);
        return -ENOMEM;
    }

    buffer = (void *)__get_free_pages(GFP_KERNEL | __GFP_ZERO, get_order(BUFFER_SIZE));
    if (!buffer) {
        kfree(pages);
        device_destroy(cls, MKDEV(major, 0));
        class_destroy(cls);
        unregister_chrdev(major, DEVICE_NAME);
        return -ENOMEM;
    }

    buffer_phys = virt_to_phys(buffer);

    for (i = 0; i < PAGE_COUNT; i++) {
        pages[i] = virt_to_page(buffer + i * PAGE_SIZE);
    }

    printk(KERN_INFO "Contig mmap module loaded. Major: %d, Phys addr: 0x%lx\n", major, buffer_phys);
    return 0;
}

static void __exit contig_mmap_exit(void) {
    device_destroy(cls, MKDEV(major, 0));
    class_destroy(cls);
    free_pages((unsigned long)buffer, get_order(BUFFER_SIZE));
    kfree(pages);
    unregister_chrdev(major, DEVICE_NAME);
    printk(KERN_INFO "Contig mmap module unloaded\n");
}

module_init(contig_mmap_init);
module_exit(contig_mmap_exit);

MODULE_LICENSE("GPL");
MODULE_DESCRIPTION("Contiguous Memory Mmap Module");
MODULE_AUTHOR("Your Name");